
import agensgraph
import psycopg2

import sys
import os

conn = psycopg2.connect('dbname=benetl host=127.0.0.1 user=lpk')
print("Connect to database successfully")


# 获取当前bulk里的所有元组的数量
def totalNumForBulk(limit,offset):
	cur = conn.cursor()
	tmpStr = "SELECT * from trainURI2id limit " + str(limit) + " offset " + str(offset)
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	return len(tmpList)

def totalNum():
	cur = conn.cursor()
	tmpStr = "SELECT * from trainURI2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	return len(tmpList)

def totalRelNum():
	cur = conn.cursor()
	tmpStr = "SELECT uri, id from relation2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	return len(tmpList)

def totalEntNum():
	cur = conn.cursor()
	tmpStr = "SELECT uri, id from entity2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	return len(tmpList)

def getList():
	cur = conn.cursor()
	tmpStr = "SELECT * from trainURI2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup)) 
	return retList 

# 获取当前bulk里的所有元组
def getListForBulk(limit,offset):
	cur = conn.cursor()
	tmpStr = "SELECT * from trainURI2id limit " + str(limit) + " offset " + str(offset)
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup))
	return retList

def getRandomList():
	cur = conn.cursor()
	tmpStr = "SELECT * from trainURI2id order by random() limit 10"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
			retList.append(list(tup))
	return retList

def getRandomList(k):
	cur = conn.cursor()
	cur.execute("SELECT * from trainURI2id order by random() limit %s",(k,))
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup))
	return retList

def getEntList():
	cur = conn.cursor()
	tmpStr = "SELECT uri, id from entity2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup)) 
	return retList

def getRelList():
	cur = conn.cursor()
	tmpStr = "SELECT uri, id from relation2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup)) 
	return retList

def getLatestEnt():
	cur = conn.cursor()
	tmpStr = "SELECT id, embedding from entity2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup)) 
	return retList

def getLatestRel():
	cur = conn.cursor()
	tmpStr = "SELECT id, embedding from relation2id"
	cur.execute(tmpStr)
	tmpList = cur.fetchall()
	retList = []
	for tup in tmpList:
		retList.append(list(tup)) 
	return retList

def writeRelationToDatabase():
	cur = conn.cursor()
	tmpStr = "TRUNCATE relation2id"
	cur.execute(tmpStr)
	tmpStr = "COPY relation2id FROM '/home/lpk/data/Test/relation2vec.txt' (DELIMITER E'\t')"
	cur.execute(tmpStr)
	conn.commit()
	return

def writeEntityToDatabase():
	cur = conn.cursor()
	tmpStr = "TRUNCATE entity2id"
	cur.execute(tmpStr)
	tmpStr = "COPY entity2id FROM '/home/lpk/data/Test/entity2vec.txt' (DELIMITER E'\t')"
	cur.execute(tmpStr)
	conn.commit()
	return

def closeConnect():
	conn.close()
	return 

# 逐条写入数据库，测试使用
def updateRelationEmbedding(value, index):
	cur = conn.cursor()
	tmpStr = "UPDATE relation2id SET embedding = ARRAY[" + value + "] WHERE id = " + str(index)
	cur.execute(tmpStr)
	conn.commit()
	return

def updateEntityEmbedding(value, index):
	# conn = psycopg2.connect('dbname=benetl host=127.0.0.1 user=lpk')
	cur = conn.cursor()
	tmpStr = "UPDATE entity2id SET embedding = ARRAY[" + value + "] WHERE id = " + str(index)
	cur.execute(tmpStr)
	conn.commit()
	return
